from concurrent.futures import ProcessPoolExecutor
from concurrent.futures import as_completed
import logging
from multiprocessing import Process
from multiprocessing import Queue
import os
import queue
from typing import Dict
from typing import List
from typing import Sequence
from typing import Set
from typing import Tuple

import numpy as np
from scipy.stats import norm
from sklearn.preprocessing import StandardScaler

from .base import Score
from .base import Scorer
from ..graph import GraphFactory
from ..model.case import Case
from ..model.case import CaseData
from ..model.graph import Graph
from ..model.graph import LoadingInvalidGraphException
from ..model.graph import Node
from ..utils import dump_json
from ..utils import load_json
from ..utils import dump_csv  # 添加 dump_csv 导入
from ..utils import load_csv  # 添加 load_csv 导入
from ..utils import require_logging
from tqdm.auto import tqdm
import time


class Evaluation:
    """
    评估结果类
    用于计算和存储根因分析模型的评估指标
    """

    def __init__(self, recommendation_num: int = 5):
        """
        初始化评估器
        Args:
            recommendation_num: 推荐列表的长度，即评估时考虑的top-k个结果，默认为5
        """
        # 存储推荐列表长度
        self._recommendation_num = recommendation_num
        # 初始化不同k值下的准确率字典
        self._accuracy = {k: 0.0 for k in range(1, recommendation_num + 1)}
        # 初始化基于实体前缀的准确率字典
        self._entity_prefix_accuracy = {
            k: 0.0 for k in range(1, recommendation_num + 1)
        }
        # 存储所有测试案例的排序结果，包括节点和分数
        self._ranks_with_scores: List[List[Tuple[Node, float]]] = []
        # 存储所有测试案例的排序结果（仅节点）
        self._ranks: List[List[Node]] = []
        # 存储答案在排名中的平均位置
        self._avg_rank_position = 0.0
        # 存储答案前缀在排名中首次出现的平均位置
        self._avg_prefix_first_position = 0.0
        # 记录案例数量
        self._case_count = 0

    def _get_entity_prefix(self, entity: str) -> str:
        """
        获取实体的前缀（& 符号前的部分）
        """
        if "&" in entity:
            return entity.split("&")[0]
        return entity

    def _check_entity_prefix_match(self, node1: Node, node2: Node) -> bool:
        """
        检查两个节点的实体前缀是否匹配
        """
        prefix1 = self._get_entity_prefix(node1.entity)
        prefix2 = self._get_entity_prefix(node2.entity)
        return prefix1 == prefix2

    def __call__(self, ranks: Sequence[Tuple[Node, Score]], answers: Set[Node]):
        """
        评估单个测试案例的结果
        Args:
            ranks: 模型预测的(节点,分数)元组列表
            answers: 真实根因节点集合
        """
        # 保存完整的排序结果，包括节点和分数
        self._ranks_with_scores.append([(node, score.key) for node, score in ranks])

        # 提取节点列表用于评估
        nodes_only = [node for node, _ in ranks]
        self._ranks.append(nodes_only)

        # 获取真实根因数量
        answer_num = len(answers)

        # 如果没有真实根因，跳过评估
        if answer_num == 0:
            self._case_count += 1
            return

        # 计算不同k值下的准确率（完全匹配）
        for k in range(1, self._recommendation_num + 1):
            # 计算前k个预测中命中的根因数量，除以总根因数量
            self._accuracy[k] += len(answers.intersection(nodes_only[:k])) / answer_num

            # 计算基于实体前缀的准确率
            prefix_matches = 0
            for answer_node in answers:
                for rank_node in nodes_only[:k]:
                    if self._check_entity_prefix_match(answer_node, rank_node):
                        prefix_matches += 1
                        break
            self._entity_prefix_accuracy[k] += prefix_matches / answer_num

        # 计算答案在排名中的平均位置
        rank_positions = []
        for answer_node in answers:
            # 查找答案节点在排名中的位置（从1开始）
            if answer_node in nodes_only:
                rank_positions.append(nodes_only.index(answer_node) + 1)
            else:
                # 如果答案不在排名中，将其位置设为排名长度+1
                rank_positions.append(len(nodes_only) + 1)

        # 计算平均排名位置
        if rank_positions:
            avg_position = sum(rank_positions) / len(rank_positions)
            self._avg_rank_position += avg_position

        # 计算答案前缀在排名中首次出现的位置
        prefix_first_positions = []
        for answer_node in answers:
            answer_prefix = self._get_entity_prefix(answer_node.entity)
            # 查找具有相同前缀的节点在排名中首次出现的位置
            prefix_found = False
            for i, rank_node in enumerate(nodes_only):
                rank_prefix = self._get_entity_prefix(rank_node.entity)
                if answer_prefix == rank_prefix:
                    prefix_first_positions.append(i + 1)  # 位置从1开始
                    prefix_found = True
                    break
            if not prefix_found:
                # 如果没有找到相同前缀的节点，位置设为排名长度+1
                prefix_first_positions.append(len(nodes_only) + 1)

        # 计算前缀首次出现的平均位置
        if prefix_first_positions:
            avg_prefix_position = sum(prefix_first_positions) / len(
                prefix_first_positions
            )
            self._avg_prefix_first_position += avg_prefix_position

        self._case_count += 1

    @property
    def num(self) -> int:
        """
        获取已评估的测试案例数量
        Returns:
            评估过的案例数量
        """
        return len(self._ranks)

    def dump(self, filename: str) -> None:
        """
        将评估结果保存到文件
        Args:
            filename: 输出文件路径
        """
        # 检查文件扩展名，确保使用CSV
        if not filename.endswith(".csv"):
            filename = os.path.splitext(filename)[0] + ".csv"
    
        csv_data = []
    
        # 如果有带分数的排名结果，使用它们
        if self._ranks_with_scores:
            # 遍历所有案例的排名结果
            for case_idx, ranks in enumerate(self._ranks_with_scores):
                # 遍历每个节点及其排名
                for rank_idx, (node, score) in enumerate(ranks):
                    # 处理不同类型的分数
                    if isinstance(score, tuple):
                        # 如果是元组，使用第一个元素作为分数
                        score_value = float(score[0])
                    elif isinstance(score, Score):
                        score_value = float(score.score)
                    else:
                        # 其他情况尝试直接转换为浮点数
                        score_value = float(score)
                    
                    # 创建一行CSV数据
                    csv_data.append([
                        case_idx,
                        rank_idx + 1,
                        node.entity,
                        node.metric,
                        score_value
                    ])
        else:
            # 如果没有带分数的排名结果，使用普通排名结果
            for case_idx, ranks in enumerate(self._ranks):
                # 遍历每个节点及其排名
                for rank_idx, node in enumerate(ranks):
                    # 创建一行CSV数据：案例索引,排名,实体,指标,分数
                    csv_data.append(
                        [
                            case_idx,
                            rank_idx + 1,
                            node.entity,
                            node.metric,
                            0.0,  # 没有分数信息，使用默认值0.0
                        ]
                    )
    
        # 按分数降序排序
        csv_data.sort(key=lambda x: x[4], reverse=True)
    
        # 保存CSV文件
        dump_csv(
            filename=filename,
            data=csv_data,
            headers=["case_id", "rank", "entity", "metric", "score"],
        )

    def load(self, filename: str, answers: Sequence[Set[Node]]) -> None:
        """
        从文件加载评估结果并计算评估指标
        Args:
            filename: 输入文件路径
            answers: 所有测试案例的真实根因集合
        """
        self._ranks = []
        self._accuracy = {k: 0.0 for k in range(1, self._recommendation_num + 1)}
        self._entity_prefix_accuracy = {
            k: 0.0 for k in range(1, self._recommendation_num + 1)
        }
        self._avg_rank_position = 0.0
        self._case_count = 0

        # 检查文件扩展名，确保使用CSV
        if not filename.endswith(".csv"):
            filename = os.path.splitext(filename)[0] + ".csv"

        if os.path.exists(filename):
            # 从CSV加载数据
            case_ranks = {}
            reader = load_csv(filename)
            next(reader)  # 跳过表头
    
            for row in reader:
                case_id, rank, entity, metric, score = row
                case_id = int(case_id)
                rank = int(rank)
                score = float(score)
    
                if case_id not in case_ranks:
                    case_ranks[case_id] = []
    
                # 创建节点对象并添加到对应案例的排名列表
                node = Node(entity=entity, metric=metric)
                case_ranks[case_id].append((rank, score, node))
    
            # 按案例ID和排名排序
            for case_id in sorted(case_ranks.keys()):
                sorted_items = sorted(case_ranks[case_id], key=lambda x: (x[0], -x[1]))
                # 创建 (node, Score) 元组列表
                sorted_pairs = [(node, Score(score)) for _, score, node in sorted_items]
                # 更新排名列表
                self._ranks.append([node for _, _, node in sorted_items])
                # 调用 __call__ 方法进行评估
                self(sorted_pairs, answers[case_id])
        else:
            # 如果CSV不存在，尝试从JSON加载（兼容旧格式）
            report: List[List[dict]] = load_json(filename)
            for ranks in report:
                self._ranks.append(
                    [
                        Node(entity=node["entity"], metric=node["metric"])
                        for node in ranks
                    ]
                )


    def accuracy(self, k: int) -> float:
        """
        计算top-k准确率(AC@k)
        对于每个案例，AC@k = 前k个预测中的正确根因数量 / 真实根因总数
        最终返回所有案例的平均AC@k

        Args:
            k: 考虑的预测结果数量
        Returns:
            平均AC@k，如果k无效或没有测试案例则返回None
        """
        if k not in self._accuracy or not self._ranks:
            return None
        return self._accuracy[k] / len(self._ranks)

    def entity_prefix_accuracy(self, k: int) -> float:
        """
        计算基于实体前缀的top-k准确率(EP-AC@k)
        对于每个案例，EP-AC@k = 前k个预测中与真实根因实体前缀匹配的数量 / 真实根因总数

        Args:
            k: 考虑的预测结果数量
        Returns:
            平均EP-AC@k，如果k无效或没有测试案例则返回None
        """
        if k not in self._entity_prefix_accuracy or not self._ranks:
            return None
        return self._entity_prefix_accuracy[k] / len(self._ranks)

    def average(self, k: int) -> float:
        """
        计算平均准确率(Avg@k)
        Avg@k = (AC@1 + AC@2 + ... + AC@k) / k

        Args:
            k: 最大考虑的预测结果数量
        Returns:
            平均准确率，如果k无效或没有测试案例则返回None
        """
        if k not in self._accuracy or not self._ranks:
            return None
        return sum(self.accuracy(i) for i in range(1, k + 1)) / k

    def entity_prefix_average(self, k: int) -> float:
        """
        计算基于实体前缀的平均准确率(EP-Avg@k)
        EP-Avg@k = (EP-AC@1 + EP-AC@2 + ... + EP-AC@k) / k

        Args:
            k: 最大考虑的预测结果数量
        Returns:
            基于实体前缀的平均准确率，如果k无效或没有测试案例则返回None
        """
        if k not in self._entity_prefix_accuracy or not self._ranks:
            return None
        return sum(self.entity_prefix_accuracy(i) for i in range(1, k + 1)) / k

    @property
    def avg_rank_position(self) -> float:
        """
        获取答案在排名中的平均位置
        Returns:
            平均排名位置，如果没有测试案例则返回None
        """
        if not self._case_count:
            return None
        return self._avg_rank_position / self._case_count

    @property
    def avg_prefix_first_position(self) -> float:
        """
        获取答案前缀在排名中首次出现的平均位置
        Returns:
            前缀首次出现的平均位置，如果没有测试案例则返回None
        """
        if not self._case_count:
            return None
        return self._avg_prefix_first_position / self._case_count
